#include "server.h"

EchoServer::EchoServer(const std::vector<int>& ports, int maxEchoCount) 
    : ports_(ports), maxEchoCount_(maxEchoCount), 
      epollFd_(-1), totalConnections_(0) {
    // 初始化echoCount数组
    memset(echoCount_, 0, sizeof(echoCount_));
}

EchoServer::~EchoServer() {
    for (int fd : listenFds_) close(fd);
    if (epollFd_ != -1) close(epollFd_);
}

void EchoServer::setNonBlocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags != -1) fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

int EchoServer::sendAll(int fd, const char* buffer, int len) {
    int total = 0;
    while (total < len) {
        int n = write(fd, buffer + total, len - total);
        if (n == -1) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) return -1;
            return -2;
        }
        total += n;
    }
    return total;
}

bool EchoServer::init() {
    epollFd_ = epoll_create1(0);
    if (epollFd_ == -1) {
        std::cerr << "Epoll create failed: " << strerror(errno) << std::endl;
        return false;
    }

    for (int port : ports_) {
        int listenFd = socket(AF_INET, SOCK_STREAM, 0);
        if (listenFd == -1) {
            std::cerr << "Create socket for port " << port << " failed: " << strerror(errno) << std::endl;
            return false;
        }

        int opt = 1;
        setsockopt(listenFd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
        // 优化：增大监听队列长度
        setsockopt(listenFd, SOL_SOCKET, SO_RCVBUF, &opt, sizeof(opt));

        struct sockaddr_in addr;
        memset(&addr, 0, sizeof(addr));
        addr.sin_family = AF_INET;
        addr.sin_addr.s_addr = INADDR_ANY;
        addr.sin_port = htons(port);
        if (bind(listenFd, (struct sockaddr*)&addr, sizeof(addr)) == -1) {
            std::cerr << "Bind port " << port << " failed: " << strerror(errno) << std::endl;
            close(listenFd);
            return false;
        }

        // 优化：使用最大监听队列长度
        if (listen(listenFd, SOMAXCONN) == -1) {
            std::cerr << "Listen port " << port << " failed: " << strerror(errno) << std::endl;
            close(listenFd);
            return false;
        }

        setNonBlocking(listenFd);
        struct epoll_event ev;
        ev.events = EPOLLIN;
        ev.data.fd = listenFd;
        if (epoll_ctl(epollFd_, EPOLL_CTL_ADD, listenFd, &ev) == -1) {
            std::cerr << "Add port " << port << " to epoll failed: " << strerror(errno) << std::endl;
            close(listenFd);
            return false;
        }

        listenFds_.push_back(listenFd);
        std::cout << "Server listening on port " << port << std::endl;
    }

    return true;
}

void EchoServer::start() {
    struct epoll_event events[MAX_EVENTS];
    char buffer[BUFFER_SIZE];

    while (true) {
        int nfds = epoll_wait(epollFd_, events, MAX_EVENTS, -1);
        if (nfds == -1) {
            if (errno == EINTR) continue;
            std::cerr << "Epoll wait failed: " << strerror(errno) << std::endl;
            break;
        }

        // 优化：批量处理事件，减少循环内开销
        for (int i = 0; i < nfds; ++i) {
            int fd = events[i].data.fd;
            uint32_t event = events[i].events;

            bool isListenFd = false;
            for (int listenFd : listenFds_) {
                if (fd == listenFd) {
                    isListenFd = true;
                    break;
                }
            }

            if (isListenFd) {
                // 批量接受新连接
                while (true) {
                    struct sockaddr_in clientAddr;
                    socklen_t clientAddrLen = sizeof(clientAddr);
                    int clientFd = accept(fd, (struct sockaddr*)&clientAddr, &clientAddrLen);
                    if (clientFd == -1) {
                        if (errno == EAGAIN || errno == EWOULDBLOCK || errno == ECONNABORTED) 
                            break;
                        std::cerr << "Accept failed: " << strerror(errno) << std::endl;
                        break;
                    }

                    // 检查fd是否超过最大限制
                    if (clientFd >= MAX_FD) {
                        std::cerr << "Client fd " << clientFd << " exceeds MAX_FD " << MAX_FD << std::endl;
                        close(clientFd);
                        continue;
                    }

                    setNonBlocking(clientFd);
                    totalConnections_++;
                    echoCount_[clientFd] = 0;

                    // 优化：减少日志输出（每100个连接打印一次）
                    if (totalConnections_ % 100 == 0) {
                        std::cout << "Total connections: " << totalConnections_ 
                                  << " (latest fd: " << clientFd << ")" << std::endl;
                    }

                    struct epoll_event ev;
                    ev.events = EPOLLIN;
                    ev.data.fd = clientFd;
                    epoll_ctl(epollFd_, EPOLL_CTL_ADD, clientFd, &ev);
                }
            } 
            else {
                if (event & EPOLLIN) {
                    int nbytes = read(fd, buffer, BUFFER_SIZE - 1);
                    if (nbytes <= 0) {
                        // 优化：仅关键日志
                        if (fd % 1000 == 0) { // 每1000个fd打印一次
                            std::cout << "Connection closed (fd: " << fd << ")" << std::endl;
                        }
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        pendingData_.erase(fd);
                        continue;
                    }

                    buffer[nbytes] = '\0';
                    echoCount_[fd]++;
                    
                    // 优化：减少日志（每100次交互打印一次）
                    if (echoCount_[fd] % 100 == 0) {
                        std::cout << "Server fd " << fd << " processed " << echoCount_[fd] << " echoes" << std::endl;
                    }

                    int sendResult = sendAll(fd, buffer, nbytes);
                    if (sendResult == nbytes) {
                        if (echoCount_[fd] >= maxEchoCount_) {
                            epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                            close(fd);
                        }
                    } 
                    else if (sendResult == -1) {
                        pendingData_[fd] = std::string(buffer, nbytes);
                        struct epoll_event ev;
                        ev.events = EPOLLIN | EPOLLOUT;
                        ev.data.fd = fd;
                        epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                    }
                    else {
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        pendingData_.erase(fd);
                    }
                }

                if (event & EPOLLOUT) {
                    auto it = pendingData_.find(fd);
                    if (it == pendingData_.end()) {
                        struct epoll_event ev;
                        ev.events = EPOLLIN;
                        ev.data.fd = fd;
                        epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                        continue;
                    }

                    std::string& data = it->second;
                    int sendResult = sendAll(fd, data.c_str(), data.size());
                    if (sendResult == data.size()) {
                        pendingData_.erase(it);
                        
                        if (echoCount_[fd] >= maxEchoCount_) {
                            epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                            close(fd);
                        } else {
                            struct epoll_event ev;
                            ev.events = EPOLLIN;
                            ev.data.fd = fd;
                            epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                        }
                    } else {
                        pendingData_.erase(it);
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                    }
                }
            }
        }
    }
}